package dan.backend
import chisel3._
import chisel3.util._
import dan.common._ 
import dan.utility._
import dan.common.Consts._ 

class DcacheReq extends CoreBundle{
    val mask = UInt(4.W)
    val addr = UInt(vaBits.W)
    val data = Bits(xBits.W)
    val uop = new UOp()
    val uncacheable = Bool()
}

class DcacheResp extends CoreBundle{
    val uop = new UOp()
    val data = Bits(xBits.W)
}

class DcacheDataReq extends CoreBundle with HasDcacheParam {
    val idx      = UInt(nIdxBits.W)
    val offset   = UInt(log2Ceil(nRowWords).W)
    val pos  = UInt(log2Ceil(nWays).W)
    val data    = UInt(xBits.W)

}

class DcacheDataResp extends CoreBundle with HasDcacheParam {
    val data = UInt(xBits.W)
}

class DcacheDataIO extends CoreBundle with HasDcacheParam {
    val req  = Input(Valid(new DcacheDataReq))
    val resp = Output(Valid(new DcacheDataResp))
}


class DcacheDataLogic extends Module with HasDcacheParam{
    val io = IO(new CoreBundle{
        val read = Vec( memWidth, (new DcacheDataIO))
        val write = (new DcacheDataIO)
    })

    // data
    val data = Module(new DcacheData)

    // Read
    for(w <- 0 until memWidth){
        data.io.read(w) <> io.read(w)
    }
    data.io.write <> io.write
}

class DcacheData extends Module with HasDcacheParam{
    val io = IO(new CoreBundle{
        // 2读口
        val read     = Vec( memWidth, new DcacheDataIO)
        // 1写口
        val write    = new DcacheDataIO
    })

    val data = SyncReadMem(nTotalWords, UInt(xBits.W))

    // reset
    val reseting = RegInit(true.B)
    val reset_1vIdx = RegInit(0.U(n1vIdxBits.W))

    when (reseting) {
        when (reset_1vIdx === (nTotalWords - 1).U) {
            reseting := false.B
        }

        // data.write(reset_1vIdx, 0.U)
        reset_1vIdx := reset_1vIdx + 1.U
    }

    // read
    val rvalid  = io.read.map( _.req.valid)
    val rreq    = io.read.map( _.req.bits)
    val ridx1v  = rreq.map(req => Cat(req.idx, req.pos, req.offset))
 
    for (w <- 0 until memWidth) {
        io.read(w).resp := 0.U.asTypeOf(Valid(new DcacheDataResp))
        io.read(w).resp.valid := RegNext(rvalid(w))
        val rdata = data.read(ridx1v(w))
        if(!FPGA_RUN)dontTouch(rdata)
        io.read(w).resp.bits.data := rdata
    }

    // write
    val wvalid  = io.write.req.valid
    val wreq    = io.write.req.bits
    val widx1v  = Cat(wreq.idx, wreq.pos, wreq.offset)

    io.write.resp := 0.U.asTypeOf(Valid(new DcacheDataResp))
    
    when (wvalid) {
        data.write(widx1v, wreq.data)
    }
    io.write.resp.valid := RegNext(wvalid)
    io.write.resp.bits.data := DontCare

    // bypass
    val bypass = Wire(Vec(memWidth, Bool()))
    if(!FPGA_RUN)dontTouch(bypass)
    for (w <- 0 until memWidth) {
        //下周期判断转发
        bypass(w) := RegNext(rvalid(w)) && RegNext(wvalid) && IsEqual( RegNext(ridx1v(w)) , RegNext(widx1v))
        when ((bypass(w))) {
            io.read(w).resp.bits.data := RegNext(wreq.data)
        }
    }

}

object WordWrite {
    def apply(req: DcacheReq, rawData: UInt) = {
        val memSize = req.uop.memSize
        val wdata = rawData.asTypeOf(Vec(4, UInt(8.W)))
        val lob = req.data(7, 0)
        val loh = req.data(15, 0)
        val low = req.data(31, 0)
        when(memSize === 0.U) {

            when(req.mask === "b0001".U) {
                wdata(0) := lob
            }
              .elsewhen(req.mask === "b0010".U) {
                  wdata(1) := lob
              }
              .elsewhen(req.mask === "b0100".U) {
                  wdata(2) := lob
              }
              .elsewhen(req.mask === "b1000".U) {
                  wdata(3) := lob
              }

        }.elsewhen(memSize === 1.U) {

            when(req.mask === "b0011".U) {
                wdata(0) := loh(7, 0)
                wdata(1) := loh(15, 8)
            }.elsewhen(req.mask === "b1100".U) {
                wdata(2) := loh(7, 0)
                wdata(3) := loh(15, 8)
            }

        }.elsewhen(memSize === 2.U) {
            when(req.mask === "b1111".U) {
                wdata(0) := low(7, 0)
                wdata(1) := low(15, 8)
                wdata(2) := low(23, 16)
                wdata(3) := low(31, 24)
            }
        }
        wdata.asUInt
    }
}

class MMIOUnit extends Module with HasDcacheParam {
    val io = IO(new CoreBundle{
        val mmioReq  = Flipped(Decoupled(new DcacheReq))
        // mmioResp is a replay request
        val mmioResp = Decoupled(new DcacheReq)

        val smar = new SMAR
        val smaw = new SMAW
    })
    if (!FPGA_RUN) dontTouch (io)

    val s_ready :: s_fetch :: s_wb :: s_resp :: Nil = Enum(4)
    val state = RegInit(s_ready)

    val mmioReq = RegInit(0.U.asTypeOf(new DcacheReq))
    if (!FPGA_RUN) dontTouch(mmioReq)

    io.mmioReq.ready := state === s_ready

    io.mmioResp.valid     := state === s_resp
    io.mmioResp.bits      := mmioReq
    io.mmioResp.bits.data := loadDataGen(
        mmioReq.addr(1, 0),
        mmioReq.data,
        mmioReq.uop.memSize,
        mmioReq.uop.memSigned
    )

    io.smar.req.arvalid     := state === s_fetch
    io.smaw.req.awvalid     := state === s_wb
    io.smaw.req.wvalid      := state === s_wb

    io.smaw.req.wstrb       := mmioReq.mask

    io.smar.req.arlen       := AXI3Parameters.MLEN1
    io.smaw.req.awlen       := AXI3Parameters.MLEN1

    io.smar.req.arburst     := AXI3Parameters.BURST_FIXED
    io.smaw.req.awburst     := AXI3Parameters.BURST_FIXED

    io.smar.req.arsize      := mmioReq.uop.memSize
    io.smaw.req.awsize      := mmioReq.uop.memSize

    io.smar.req.araddr      := mmioReq.addr
    io.smaw.req.awaddr      := mmioReq.addr

    io.smaw.req.wdata       :=
        Mux(mmioReq.uop.memSize === 0.U, Fill(4, mmioReq.data( 7, 0)),
        Mux(mmioReq.uop.memSize === 1.U, Fill(2, mmioReq.data(15, 0)),
                                          mmioReq.data))

    io.smaw.req.wlast        := state === s_wb

    when (state === s_ready) {
        when (io.mmioReq.fire) {
            state := Mux(isStore(io.mmioReq.bits), s_wb, s_fetch)
            mmioReq := io.mmioReq.bits
        }
    } .elsewhen (state === s_fetch) {
        state := Mux(io.smar.resp.rvalid, s_resp, s_fetch)
        mmioReq.data := io.smar.resp.rdata
    } .elsewhen (state === s_wb){
        if(!FPGA_RUN){
            val debug_lo_byte = mmioReq.data(7, 0)
            val debug_lo_half = mmioReq.data(15, 0)
            dontTouch(debug_lo_byte)
            dontTouch(debug_lo_half)
        }


        state := Mux(io.smaw.resp.wready, s_resp, s_wb)
    } .elsewhen(state === s_resp) {
        state := Mux(io.mmioResp.fire, s_ready, s_resp)
    }
}

class Missarbiter extends CoreModule with HasDcacheParam {
    val io = IO(new Bundle{
        val req   = Input(Vec(memWidth, new DcacheReq))
        val alive = Input(Vec(memWidth, Bool()))
        val miss  = Input(Vec(memWidth, Bool()))

        val sendNack         = Output(Vec(memWidth , Bool()))
        val sendResp         = Output(Vec(memWidth , Bool()))
        val store_enter_mshr = Output(Bool())
        val storeFailed      = Output(Bool())

        val mshrReq = Decoupled(new DcacheReq)
    })
    if (!FPGA_RUN) dontTouch(io)

    io.sendResp    := 0.U.asTypeOf(Vec(memWidth, Bool()))
    io.sendNack    := 0.U.asTypeOf(Vec(memWidth, Bool()))
    io.storeFailed := false.B

    io.mshrReq.valid := false.B
    io.mshrReq.bits  := DontCare

    val store_enter_mshr = WireInit(false.B)
    io.store_enter_mshr := store_enter_mshr

    val wantaccess = Wire(Vec(memWidth, Bool()))
    for (i <- 0 until memWidth) {
        wantaccess(i) := io.miss(i) && io.alive(i)
    }

    when (wantaccess(0) && !wantaccess(1)) {
        // 上个周期成功进入了一条，这个周期又一个miss_store不能存进去（由于hasStore被RegNext，ready来不及变化）相当于做hasStore的转发
        // 不能存进去，同时要nack和storeFailed，通过valid为假导致的握手失败，来实现
        val do_st_as_miss = isStore(io.req(0)) && RegNext(store_enter_mshr)
        io.sendResp(0) := false.B
        io.sendNack(0) := !io.mshrReq.fire
        io.storeFailed := !io.mshrReq.fire && isStore(io.req(0))

        store_enter_mshr := io.mshrReq.fire && isStore(io.req(0))
        
        io.mshrReq.valid := true.B && !do_st_as_miss
        io.mshrReq.bits  := io.req(0)
    } .elsewhen(!wantaccess(0) && wantaccess(1)) {
        io.sendResp(1) := false.B
        io.sendNack(1) := !io.mshrReq.ready

        io.mshrReq.valid := true.B
        io.mshrReq.bits  := io.req(1)
    } .elsewhen(wantaccess(0) && wantaccess(1)) { // priority: 1 > 0
        // pipeline 0
        io.sendResp(0) := false.B
        io.sendNack(0) := true.B
        io.storeFailed := isStore(io.req(0))

        // pipeline 1
        io.sendResp(1) := false.B
        io.sendNack(1) := !io.mshrReq.ready

        // send request to mshr
        io.mshrReq.valid := true.B
        io.mshrReq.bits  := io.req(1)
    }
}
